package info.wangyuan.agent.utils;

import cn.hutool.core.collection.CollUtil;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import info.wangyuan.agent.common.Constant.FileConstant;
import info.wangyuan.agent.common.Enum.FileEventType;
import info.wangyuan.agent.entity.po.FileInfo;
import info.wangyuan.agent.entity.po.KnowledgeBase;
import info.wangyuan.agent.entity.po.KnowledgeSource;
import info.wangyuan.agent.mapper.DocumentUnitMapper;
import info.wangyuan.agent.mapper.FileInfoMapper;
import info.wangyuan.agent.mapper.KnowledgeBaseMapper;
import info.wangyuan.agent.mapper.KnowledgeSourceMapper;
import info.wangyuan.agent.service.event.RagFileEvent;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

@Slf4j
@Component
public class RagRebuildService {

    private final KnowledgeBaseMapper kbMapper;
    private final KnowledgeSourceMapper ksMapper;
    private final FileInfoMapper fileMapper;
    private final DocumentUnitMapper documentUnitMapper;
    private final EmbeddingStore<TextSegment> embeddingStore;
    private final ApplicationEventPublisher publisher;

    public RagRebuildService(KnowledgeBaseMapper kbMapper, KnowledgeSourceMapper ksMapper, FileInfoMapper fileMapper, DocumentUnitMapper documentUnitMapper, EmbeddingStore<TextSegment> embeddingStore, ApplicationEventPublisher publisher) {
        this.kbMapper = kbMapper;
        this.ksMapper = ksMapper;
        this.fileMapper = fileMapper;
        this.documentUnitMapper = documentUnitMapper;
        this.embeddingStore = embeddingStore;
        this.publisher = publisher;
    }

    /**
     * 重建知识库文件及向量数据核心逻辑
     */
    @Transactional
    public void rebuild(String rootPath) throws IOException {
        log.info("========== 开始重建知识库与向量索引 ==========");

        // 清空数据库和向量数据
        clearDatabasesAndVectors();

        // 获取知识库集合
        File root = new File(rootPath);
        List<KnowledgeBase> kbList = scanKnowledgeBase(root);
        if (CollUtil.isEmpty(kbList)) {
            log.warn("未发现知识库，重建结束。");
            return;
        }

        List<Integer> fileIdList = new ArrayList<>();

        for (KnowledgeBase kb : kbList) {
            // 知识库入库
            kbMapper.insert(kb);

            // 知识源入库
            bindAgentsToKnowledgeBase(kb);

            // 获取知识库文件集合
            List<File> fileList = collectFiles(kb.getPath());

            // 转为文件类集合
            List<FileInfo> fileInfoList = buildFileInfoList(fileList, kb);

            // 知识库文件批量入库
            if (CollUtil.isNotEmpty(fileInfoList)) {
                fileMapper.insert(fileInfoList);
                fileIdList.addAll(
                        fileInfoList.stream()
                                .map(FileInfo::getId)
                                .filter(Objects::nonNull)
                                .toList()
                );
            }
        }

        // 知识库文件切分和向量化
        if (CollUtil.isNotEmpty(fileIdList)) {
            publisher.publishEvent(new RagFileEvent<>(FileEventType.CHUNK, fileIdList));
        }

        log.info("========== 重建完成：知识库 {} 个，文件 {} 个 ==========", kbList.size(), fileIdList.size());
    }

    /**
     * 绑定 Agent 与知识库
     */
    private void bindAgentsToKnowledgeBase(KnowledgeBase kb) throws IOException {
        // 知识库元数据处理
        Set<Integer> agentIds = FolderMetaUtils.getAgentIdsSet(Path.of(kb.getPath()));
        if (CollUtil.isEmpty(agentIds)) return;

        List<KnowledgeSource> knowledgeSources = agentIds.stream()
                .map(agentId -> KnowledgeSource.builder()
                        .id(agentId)
                        .agentId(agentId)
                        .knowledgeBaseId(kb.getId())
                        .build())
                .filter(Objects::nonNull)
                .toList();

        if (CollUtil.isNotEmpty(knowledgeSources))
            ksMapper.insert(knowledgeSources);
    }

    /**
     * 清空数据库和向量数据
     */
    private void clearDatabasesAndVectors() {
        kbMapper.truncateKnowledgeBase();           // 清空知识库表
        ksMapper.truncateKnowledgeSource();         // 清空知识源表
        documentUnitMapper.truncateDocumentUnit();  // 清空文档单元表
        fileMapper.truncateFileInfo();              // 清空文件信息表
        embeddingStore.removeAll();                 // 清空向量表
    }

    /**
     * 构建 FileInfo 列表
     */
    private List<FileInfo> buildFileInfoList(List<File> fileList, KnowledgeBase kb) {
        return fileList.stream()
                .map(file -> buildFileInfo(file, kb))
                .filter(Objects::nonNull)
                .toList();
    }

    /**
     * 从指定目录递归收集指定后缀名的文件
     *
     * @param dirPath 文件目录
     * @return List<File>
     */
    private List<File> collectFiles(String dirPath) throws IOException {
        try (Stream<Path> paths = Files.walk(Paths.get(dirPath))) {
            return paths
                    .filter(Files::isRegularFile) // 只选文件
                    .filter(path -> {
                        String fileName = path.getFileName().toString().toLowerCase();
                        return FileConstant.PATH_MATCHER_EXTENSIONS.stream().anyMatch(fileName::endsWith);
                    })
                    .map(Path::toFile)
                    .collect(Collectors.toList());
        }
    }

    /**
     * 构建 FileInfo
     */
    private FileInfo buildFileInfo(File file, KnowledgeBase kb) {
        return FileInfo.builder()
                .name(file.getName())
                .knowledgeBaseId(kb.getId())
                .path(file.getAbsolutePath())
                .size(file.length())
                .suffix(getFileSuffix(file))
                .build();
    }

    /**
     * 获取文件后缀
     */
    private static String getFileSuffix(File file) {
        String name = file.getName();
        int idx = name.lastIndexOf('.');
        return idx > -1 ? name.substring(idx + 1) : "";
    }

    /**
     * 获取文件绝对路径
     */
    private static String getCanonicalPath(File file) {
        try {
            return file.getCanonicalPath();
        } catch (Exception e) {
            return file.getAbsolutePath();
        }
    }

    /**
     * 扫描知识库目录
     */
    public List<KnowledgeBase> scanKnowledgeBase(File root) {
        if (!root.exists() || !root.isDirectory()) {
            log.warn("ROOT_PATH [{}] 不存在或不是目录，跳过扫描。", getCanonicalPath(root));
            return Collections.emptyList();
        }

        File[] dirs = root.listFiles(File::isDirectory);
        if (dirs == null || dirs.length == 0) {
            log.warn("ROOT_PATH [{}] 下未发现知识库目录", getCanonicalPath(root));
            return Collections.emptyList();
        }

        List<KnowledgeBase> kbList = new ArrayList<>(dirs.length);
        for (File dir : dirs) {
            KnowledgeBase kb = KnowledgeBase.builder()
                    .name(dir.getName())
                    .path(getCanonicalPath(dir))
                    .build();
            kbList.add(kb);
            log.info("发现知识库：{}", kb.getName());
        }

        return kbList;
    }
}
